import torch
import torch.nn as nn
import torchvision
from torch.utils.data import DataLoader
from tqdm import tqdm
import time

# 检查 GPU 是否可用并打印信息
def check_device():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    return device

class Wang(nn.Module):
    def __init__(self):
        super(Wang, self).__init__()
        self.model1 = nn.Sequential(
            nn.Conv2d(3, 32, 5, padding=2, stride=1),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, padding=2, stride=1),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, padding=2, stride=1),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 64, 5, padding=2, stride=1),  # 新增卷积层
            nn.MaxPool2d(2),
            nn.Conv2d(64, 64, 5, padding=2, stride=1),  # 再新增卷积层
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*1*1, 64),  # 调整输入维度
            nn.Linear(64, 10)
        )

    def forward(self, x):
        return self.model1(x)

# 训练函数
def train(dataloader, model, loss_fn, optimizer, epochs):
    start_time = time.time()  # 记录训练开始时间
    for epoch in range(epochs):
        print(f"Epoch {epoch+1}\n-------------------------------")
        model.train()
        epoch_loss = 0
        for imgs, targets in tqdm(dataloader, desc=f"Epoch {epoch+1}"):
            imgs, targets = imgs.to(device), targets.to(device)  # 将数据移动到 GPU
            output = model(imgs)
            loss = loss_fn(output, targets)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            epoch_loss += loss.item()
        avg_loss = epoch_loss / len(dataloader)
        print(f"Average Loss: {avg_loss:>8f}")
    end_time = time.time()  # 记录训练结束时间
    total_time_seconds = end_time - start_time
    total_time_minutes = total_time_seconds / 60
    print(f"Total Training Time: {total_time_seconds:.2f} seconds")
    print(f"Total Training Time: {total_time_minutes:.2f} minutes")

# 主程序入口
if __name__ == '__main__':
    device = check_device()  # 检查设备并打印信息
    # 加载测试集
    dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=torchvision.transforms.ToTensor())
    dataloader = DataLoader(dataset, batch_size=128, shuffle=True, num_workers=4)  # 增加 batch_size 和 num_workers

    loss_fn = nn.CrossEntropyLoss()
    wang = Wang().to(device)  # 将模型移动到 GPU
    optim = torch.optim.SGD(wang.parameters(), lr=0.001)

    epochs = 10
    train(dataloader, wang, loss_fn, optim, epochs)
    print("Done!")
